package com.chujian.framework.util;

import cn.hutool.core.collection.CollUtil;
import com.chujian.framework.core.base.BaseTreeEntity;

import java.io.Serializable;
import java.util.*;
import java.util.stream.Collectors;

/**
 * 树形结构工具类
 *
 * @author chujian
 */
public class TreeUs {

    private TreeUs() {
    }

    /**
     * 通用的tree构建
     *
     * @param items       items
     * @param topId       父级ID
     * @param treeCompose tree构建器
     * @param <T>         T
     * @return List<T>
     */
    public static <T> List<T> toTree(List<T> items, Serializable topId, ITreeCallback<T> treeCompose) {
        if (CollUtil.isEmpty(items)) return items;

        // 收集
        Map<Serializable, T> map = new HashMap<>();
        for (T item : items) {
            Serializable id = treeCompose.getId(item);
            map.put(id, item);
        }
        // 组合
        for (T item : items) {
            T pItem = map.get(treeCompose.getPId(item));
            if (null != pItem) {
                List<T> children = treeCompose.getChildren(pItem);
                if (null != children) {
                    children.add(item);
                    treeCompose.afterAddChild(pItem, item);
                }
            }
        }
        // 返回
        return items
                .stream()
                .filter(e -> topId.equals(treeCompose.getPId(e)))
                .collect(Collectors.toList());
    }

    /**
     * TreeEntity 的 toTree
     *
     * @param items items
     * @param <T>   T
     * @return List<T>
     */
    public static <T extends BaseTreeEntity<T>> List<T> toEntityTree(List<T> items) {
        if (CollUtil.isEmpty(items)) return items;

        // 收集
        Map<Serializable, T> map = new HashMap<>();
        for (T item : items) {
            Serializable id = item.getId();
            map.put(id, item);
        }

        // 组合
        Set<Serializable> sonIds = new HashSet<>(); // 收集子节点（非父节点）
        for (T item : items) {
            Serializable parentId = item.getParentId();
            T parent = map.get(parentId);
            if (null != parent) {
                parent.initChildren();
                parent.addChildren(item);

                // 收集非父节点
                sonIds.add(item.getId());
            }
        }

        // 返回
        return items
                .stream()
                .filter(e -> !sonIds.contains(e.getId()))
                .collect(Collectors.toList());
    }

    /**
     * tree回调
     *
     * @param <T>
     */
    public interface ITreeCallback<T> {

        Serializable getId(T item);

        Serializable getPId(T item);

        List<T> getChildren(T item);

        default void afterAddChild(T parent, T item) {

        }
    }
}
